from datetime import date, timedelta
from pickle import dump
from h5py import File
from numpy import zeros, sum
from pandas import DataFrame, pivot_table
from stock import Stock, StockCollection
from pdb import set_trace

class StockHolding(Stock):
    """
    Class to describe a holding of a market stock. Adds to Stock object
    quantity - current quantity of shares held
    avg_cost - average cost of the current shares held
    total_cost - total cost of shares (using the average cost)
    TODO eliminate this class
    """
    def __init__(self, ticker, quantity, price):
        super(StockHolding,self).__init__(ticker)
        self.quantity = quantity
        self.avg_price = price
        self.total_cost = quantity * price

    def add(self, quantity, price):
        """
        Add shares to the holding
        quantity - number of shares to add
        price - cost of the shares
        """
        self.quantity += quantity
        self.total_cost += price * quantity
        self.avg_price = self.total_cost / self.quantity

    def remove(self, quantity):
        """
        Remove shares from the holding
        quantity - number of shares to remove
        """
        self.quantity -= quantity
        self.total_cost -= self.avg_price * quantity

class TransactionError(Exception):
    def __init__(self, value):
         self.value = value
    def __str__(self):
         return repr(self.value)

class Portfolio(object):
    """
    Base class to contain all asset
    holdings
    """
    def __init__(self):
        self.cash = 0                       # float of cash in market sweep
        self.holdings = {}                  # dictionary of holdings key is ticker
        self.translog = []                  # list of dictionary of transcations
        self.historical_quotes = None       # numpy structured array of historical quote data (Note not saved to hdf5 file)

    @classmethod
    def from_log_array(cls, transaction_log):
        """
        Create a portfolio object from a transaction log.
        transaction_log - dictionary of transactions from
                which to create the portfolio object
        return instance of a Portfolio object
        """
        transaction_log = sorted(transaction_log, key=lambda x: x['date'])
        port = cls()
        for t in transaction_log:
            if t['type'] == 'deposit':
                port.deposit(t['quantity'], t['date'])
            elif t['type'] == 'buy':
                port.buy(t['ticker'], t['date'], t['price'], 
                        t['quantity'])
            elif t['type'] == 'sell':
                port.sell(t['ticker'], t['date'], t['price'], 
                        t['quantity'])
        return port

    @staticmethod
    def load_log_array(fname, hdf_path):
        """
        Load a transaction log from an hdf5 file
        fname - string name of hdf5 file name
        hdf_path - string path of the hdf5 dataset of the translog
        return numpy array of the transaction log
        """
        with File(fname, mode='r') as f_h5py:
            alog = f_h5py[hdf_path].value

        temp = zeros(len(alog), dtype=[('type','|S9'), ('date',date), 
            ('ticker','|S9'), ('price','<f4'), ('quantity','<f4')])
        temp['date'] = map(date.fromordinal, alog['date'])
        for k in ['type', 'ticker', 'price', 'quantity']:
            temp[k] = alog[k]

        return temp

    def buy(self, ticker, pdate, pprice, pquantity):
        """
        Add a holding to the portfolio
        """
        if self.cash >= pprice * pquantity:
            if self.holdings.has_key(ticker):
                self.holdings[ticker] =+ pquantity
            else:
                self.holdings.update({ticker : pquantity})
            self.cash -= pprice * pquantity
            self.translog.append({'ticker':ticker, 'date':pdate, 'price':pprice, 
                'quantity':abs(pquantity), 'type':'buy'})
        else:
            raise TransactionError("Insufficient funds for puchase")

    def sell(self, ticker, sdate, sprice, squantity):
        """
        Remove a holding from the portfolio
        """
        if self.holdings.has_key(ticker):
            if self.holdings[ticker] >= squantity:
                self.holdings[ticker] -= squantity
                self.cash += sprice * squantity
                self.translog.append({'ticker':ticker, 'date':sdate, 'price':sprice, 
                    'quantity':-1*abs(squantity), 'type':'sell'})
                if self.holdings[ticker] == 0:
                    del self.holdings[ticker]
            else:
                raise TransactionError('Insufficient quantity of %s to complete transaction'%ticker)
        else:
            raise TransactionError('No shares of %s in portfolio'%ticker)

    def deposit(self, cash, d):
        """
        Deposit cach into the portfolio account
        """
        self.cash += cash
        self.translog.append({'ticker':'cash', 'date':d, 'price':1, 
            'quantity':cash, 'type':'deposit'})

    def withdrawl(self, cash, d):
        """
        Deposit cach into the portfolio account
        """
        self.cash -= cash
        self.translog.append({'ticker':'cash', 'date':d, 'price':1, 
            'quantity':cash, 'type':'withdrawl'})

    def get_transactions(self, ticker=None, startdate=None, enddate=None):
        """
        Get transactions of ticker from the time interval startdate to enddate.
        ticker - string
        startdate - date object
        enddate - date object
        return dictionary of selected transactions
        """
        if startdate is None:
            startdate = self.translog[0]['date']
        if enddate is None:
            enddate = self.translog[-1]['date']

        select = [t for t in self.translog if t['date']>startdate and t['date']<enddate]
        return select

    def get_holdings(self, date=None):
        """
        Get the portfolio holdings on a specified date.
        date - date at which to determine holdings. If date is
                None then return holdings at the most recent time
        return dictionary of holdings
        """
        if date is None:
            return self.holdings

        translog = DataFrame(self.translog)[:date]
        pvt = pivot_table(translog, values=['quantity'], rows=['date'], cols=['ticker'], 
                fill_value=0, aggfunc={'quantity':sum})
        holdings = sum(pvt, axis=0)
        return holdings

        if date is None:
            for k, v in self.holdings.iteritems():
                holdings[k] = v.quantity
            return holdings
        else:
            trans = self.get_transactions(ticker=None, startdate=None, enddate=date)
            trans = sorted(trans, key=lambda x: x['date'])
            for t in trans:
                if t['type'] == 'buy':
                    if holdings.has_key(t['ticker']):
                        holdings[t['ticker']] += t['quantity']
                    else:
                        holdings[t['ticker']] = t['quantity']
                elif t['type'] == 'sell':
                    holdings[t['ticker']] -= t['quantity']
                    if holdings[t['ticker']] == 0:
                        del holdings[t['ticker']]
                    elif holdings[t['ticker']] < 0:
                        raise('Negative quantity encountered')
            return holdings

    def get_cash(self, vdate):
        """
        Calculate the portfolio cash value at date vdate
        return float of cash value
        """
        cash = 0
        if vdate is None:
            return self.cash
        else:
            trans = self.get_transactions(ticker=None, startdate=None, enddate=vdate)
            trans = sorted(trans, key=lambda x: x['date'])
            for t in trans:
                if t['type'] == 'deposit':
                    cash += t['quantity']
                elif t['type'] == 'withdrawl':
                    cash -= t['quantity']
                if t['type'] == 'buy':
                    cash -= t['quantity'] * t['price']
                elif t['type'] == 'sell':
                    cash += t['quantity'] * t['price']
                if cash < 0:
                    raise('Negative cash encountered')
            return cash

    def __create_stock_collection(self):
        """
        Create a StockCollection from all assets involved in any transactions,
        and import historicatl quote data from the first transaction date to
        the last transaction date.
        return : StockCollection object
        """
        tickers = list(set([t['ticker'] for t in self.translog]))
        stkcol = StockCollection(tickers)
        stkcol.import_historical_data(self.translog[0]['date'], self.translog[-1]['date'])
        self.historical_quotes = stkcol.to_array()
    
    def NPV(self):
        """
        Calculate the NPV of the portfolio.
        """
        NPV = self.cash

        for tck, stk in self.holdings.iteritems():
            stk.import_quote()
            price = stk.quote['price']
            NPV += stk.quantity * price
        return NPV

    def get_value(self, vdate):
        """
        Calculate the value of the portfolio at the close on
        a given date, vdate.
        vdate - datetime date object
        return float value of portfolio
        TODO
        """
        if self.historical_quotes == None:
            market = self.__create_stock_collection()

        asset_value = 0
        for t, q in self.holdings.iteritems():
            asset_value += market[t].get_quotes(vdate)['close'] * q

        cash_value = self.get_cash(vdate)
        
        value = {'cash':cash_value, 'assets':asset_value}
        return value
        
    def running_value(self):
        """
        Calculate the running value of the portfolio
        return : numpy array of [date, value]
        """
    #    if self.historical_quotes == None:
    #        self.__create_stock_collection()

        translog = DataFrame(self.translog)
        pvt = pivot_table(translog, values=['quantity', 'type'], rows=['date'], cols=['ticker'], 
                fill_value=0, aggfunc={'quantity':sum, 'type':lambda x: str(x).find('sell')!=-1})
        quantity = pvt['quantity']
        sell = pvt['type']
        
        sell[sell==True] = -1
        sell[(sell==False) or (sell==0)] = 1
        quantity = quantity * sell
        return sum(quantity, axis=0)

        assets = {t['ticker']:0 for t in self.translog}
        assets.update({'cash':0, 'deposit':0, 'date':self.translog[0]['date']})
        holdings = [assets]

        for i, t in enumerate(self.translog):
            if t['type'] == 'buy':
                assets[t['ticker']] += t['quantity']
                assets['cash'] -= t['price']*t['quantity']
            elif t['type'] == 'sell':
                assets[t['ticker']] -= t['quantity']
                assets['cash'] +=t ['price']*t['quantity']
            elif t['type'] == 'deposit':
                assets['deposit'] += t['quantity']
                assets['cash'] += t['quantity']

            if i+1 < len(self.translog):
                if self.translog[i+1]['date'] > t['date']:
                    assets['date'] = t['date']
                    holdings.append(assets.copy())
            else:
                assets['date'] = t['date']
                holdings.append(assets.copy())

        holdings = DataFrame(holdings)

        etf_quant = holdings[assets.keys()]

        return holdings

    def to_array(self):
        """
        Convert the list of dictionaries of the transaction log
        to a structured array.
        return numpy structured array
        """
        n = len(self.translog)
        alog = zeros((n), dtype=[('type','|S9'), ('date',date), 
            ('ticker','|S9'), ('price','<f4'), ('quantity','<f4'), 
            ('gain', '<f4')])
        for i, t in enumerate(self.translog):
            for k in ['type', 'date', 'ticker', 'price', 'quantity']:
                alog[k][i] = t[k]
                if t['type'] == 'sell':
                    alog['gain'][i] = t['gain']
                elif t['type'] == 'buy':
                    alog['gain'][i] = 0
        return alog

    def save_log(self, fname, h5py_path):
        """
        Save the portfolio transaction log to an hdf5 file
        that can be used to subsequently recreate the portfolio
        object.
        fname - string h5py File name
        h5py_path - string path name where to create the hdf5 file
        """
        alog = self.to_array()
        """ A new array has to be constructed because date objects 
        are not supported in hdf5. Possibly the view method could
        be used."""
        temp = zeros(len(alog), dtype=[('type','|S9'), ('date','int32'), 
            ('ticker','|S9'), ('price','<f4'), ('quantity','<f4')])
        temp['date'] = map(date.toordinal, alog['date'])
        for k in ['type', 'ticker', 'price', 'quantity']:
            temp[k] = alog[k]

        with File(fname,mode='a') as f_h5py:
            f_h5py.create_dataset(h5py_path, data=temp)

if __name__ == "__main__":
    p = Portfolio()

    startdate = date(2012,7,23)
    enddate = date(2013,7,22)

    p.deposit(2000, startdate)
    print "Starting NPV: %d \n" %p.NPV()

    purchdate = startdate + timedelta(days=7)
    # total = $1500
    buys = [('mmm',10,100), ('ge',10,20), ('pfe',10,30)]
    for ticker, n_purchase, purchprice in buys:
        p.buy(ticker, purchdate, purchprice, n_purchase)

    print "NPV after buy: %d" %p.NPV()      
    print p.get_holdings()
    print "cash: %d \n" %p.cash

    selldate = purchdate + timedelta(days=7)
    # total = $1150
    sells = [('mmm',5,110), ('ge',5,40), ('pfe',10,40)]
    for ticker, n_sell, sellprice in sells:
        p.sell(ticker, selldate, sellprice, n_sell)

    print "NPV after sell: %d" %p.NPV()      
    print p.get_holdings()
    # Should be 2000 - 1500 + 1150 = $1650
    print "cash: %d \n" %p.cash
    set_trace()
